package A一周刷爆LeetCode.A基础班.G图;

import java.util.*;

/**
 * @author wei.zhao
 * @description: 最小生成树P算法
 * @date 2022年03月08日 16:24
 */
public class Code05_Prim {

    public static Set<Edge> primMST(Graph graph) {
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(Comparator.comparing(e -> e.weight));
        // 已经连通的节点
        HashSet<Node> set = new HashSet<>();
        HashSet<Edge> result = new HashSet<>();
        for (Node node : graph.nodes.values()) {
            if (!set.contains(node)) {
                set.add(node);
                priorityQueue.addAll(node.edges);
                while (!priorityQueue.isEmpty()) {
                    Edge edge = priorityQueue.poll();
                    Node to = edge.to;
                    if (!set.contains(edge.to)) {
                        result.add(edge);
                        set.add(to);
                        priorityQueue.addAll(edge.to.edges);
                    }
                }

            }
        }
        return result;
    }

}
